package org.example;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import io.reactivex.Flowable;
import io.reactivex.schedulers.Schedulers;

import java.util.ArrayList;
import java.util.List;
import java.util.Scanner;
import java.util.concurrent.CountDownLatch;

import org.example.utils.COMMON;

public class StreamingConversation {
    public static List<Message> messages = new ArrayList<>();

    public static Integer token = 0;

    public static GenerationParam createGenerationParam() {
        return GenerationParam.builder()
                .apiKey(COMMON.API_KEY)
                .model(COMMON.MODEL)
                .messages(messages)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true) // 开启增量输出，流式返回
                .build();
    }

    public static GenerationResult callGenerationWithMessages(GenerationParam param) throws ApiException, NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        return gen.call(param);
    }

    private static Message createMessage(Role role, String content) {
        return Message.builder().role(role.getValue()).content(content).build();
    }

    public static void main(String[] args) {
        try {
            Scanner scanner = new Scanner(System.in);
            Integer times = 1;
            System.out.println("欢迎来到对话！(若需退出请输入exit)");
            messages.add(createMessage(Role.SYSTEM, "you are a helpful assistant"));
            while (true) {

                Generation gen = new Generation();
                CountDownLatch latch = new CountDownLatch(1);


                System.out.println("第" + times + "轮（当前消耗总token:" + token + "）");
                System.out.print("用户：");
                String input = scanner.nextLine();
                if (input.equals("exit")) {
                    break;
                }
                messages.add(createMessage(Role.USER, input));

                // 开始对话
                GenerationParam param = createGenerationParam();
                Flowable<GenerationResult> result = gen.streamCall(param);
                StringBuilder fullContent = new StringBuilder();
                System.out.print("AI: ");
                result
                        .subscribeOn(Schedulers.io())
                        .observeOn(Schedulers.computation())
                        .subscribe( // 处理返回的每一块内容
                                message -> {
                                    String content = message.getOutput().getChoices().get(0).getMessage().getContent();
                                    String finishReason = message.getOutput().getChoices().get(0).getFinishReason();
                                    System.out.print(content);
                                    fullContent.append(content);
                                    if (finishReason != null && !"null".equals(finishReason)) {
                                        // 单论对话结束后计算已消耗token
                                        token += message.getUsage().getTotalTokens();
                                        // 将ai的回复添加到messages中
                                        messages.add(createMessage(Role.ASSISTANT, fullContent.toString()));
                                    }
                                    System.out.flush();
                                },
                                error -> {
                                    System.err.println("\n请求失败: " + error.getMessage());
                                    latch.countDown();
                                },
                                () -> {
                                    System.out.println();
                                    latch.countDown();
                                }
                        );
                times++;
                // 主线程等待异步任务完成
                latch.await();
            }
        } catch (Exception e) {
            System.err.println("请求异常: " + e.getMessage());
            e.printStackTrace();

        }
    }
}